import os
from dotenv import load_dotenv 
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.tools import tool
from pydantic import BaseModel, Field
import requests, json
import sqlite3
import json as json_lib
from datetime import datetime
from typing import Dict, List, Any
from langchain_deepseek import ChatDeepSeek

# 加载环境变量
load_dotenv(override=True)

# 自定义SQLite持久化存储类
class CustomSQLiteMemory:
    def __init__(self, db_path: str = "chatbot_memory.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """初始化数据库表"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建对话历史表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS conversations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                thread_id TEXT NOT NULL,
                message_type TEXT NOT NULL,
                content TEXT NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                metadata TEXT
            )
        ''')
        
        # 创建索引
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_thread_id ON conversations(thread_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_timestamp ON conversations(timestamp)')
        
        conn.commit()
        conn.close()
        print(f"✅ 自定义SQLite记忆数据库已初始化: {self.db_path}")
    
    def save_message(self, thread_id: str, message_type: str, content: str, metadata: Dict = None):
        """保存消息到数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        metadata_json = json_lib.dumps(metadata) if metadata else None
        
        cursor.execute('''
            INSERT INTO conversations (thread_id, message_type, content, metadata)
            VALUES (?, ?, ?, ?)
        ''', (thread_id, message_type, content, metadata_json))
        
        conn.commit()
        conn.close()
        print(f"💾 保存消息: {thread_id} - {message_type}")
    
    def get_conversation_history(self, thread_id: str, limit: int = 50) -> List[Dict]:
        """获取对话历史"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT message_type, content, timestamp, metadata
            FROM conversations 
            WHERE thread_id = ?
            ORDER BY timestamp ASC
            LIMIT ?
        ''', (thread_id, limit))
        
        rows = cursor.fetchall()
        conn.close()
        
        history = []
        for row in rows:
            message_type, content, timestamp, metadata = row
            history.append({
                'type': message_type,
                'content': content,
                'timestamp': timestamp,
                'metadata': json_lib.loads(metadata) if metadata else None
            })
        
        print(f"📚 获取对话历史: {thread_id} - {len(history)} 条消息")
        return history
    
    def clear_conversation(self, thread_id: str):
        """清空指定对话"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('DELETE FROM conversations WHERE thread_id = ?', (thread_id,))
        deleted_count = cursor.rowcount
        
        conn.commit()
        conn.close()
        print(f"🗑️ 清空对话: {thread_id} - 删除了 {deleted_count} 条消息")

# 创建全局记忆存储实例
persistent_memory = CustomSQLiteMemory()

# 内置搜索工具
# search_tool = TavilySearch(max_results=5, topic="general")

class WeatherQuery(BaseModel):
    loc: str = Field(description="The location name of the city")

@tool(args_schema = WeatherQuery)
def get_weather(loc):
    """
    查询即时天气函数
    :param loc: 必要参数，字符串类型，用于表示查询天气的具体城市名称，\
    支持中文和英文城市名称，例如：'北京'、'Beijing'、'上海'、'Shanghai'等；
    :return：wttr.in免费天气API查询结果，返回JSON格式的天气信息，包含温度、湿度、风速、天气状况等详细信息
    """
    try:
        # Step 1.构建请求URL - 使用wttr.in免费API
        # 使用format=j1获取JSON格式数据，lang=zh获取中文输出
        url = f"https://wttr.in/{loc}?format=j1&lang=zh"
        
        # Step 2.设置请求头，模拟浏览器访问
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        
        # Step 3.发送GET请求
        response = requests.get(url, headers=headers, timeout=10)
        
        # Step 4.检查响应状态
        if response.status_code == 200:
            data = response.json()
            
            # 提取主要天气信息并格式化
            current = data.get('current_condition', [{}])[0]
            location = data.get('nearest_area', [{}])[0]
            
            # 构建简化的天气信息
            weather_info = {
                "location": {
                    "name": location.get('areaName', [{}])[0].get('value', loc),
                    "country": location.get('country', [{}])[0].get('value', ''),
                    "region": location.get('region', [{}])[0].get('value', '')
                },
                "current": {
                    "temperature": current.get('temp_C', 'N/A'),
                    "feels_like": current.get('FeelsLikeC', 'N/A'),
                    "humidity": current.get('humidity', 'N/A'),
                    "description": current.get('lang_zh', [{}])[0].get('value', current.get('weatherDesc', [{}])[0].get('value', '')),
                    "wind_speed": current.get('windspeedKmph', 'N/A'),
                    "wind_direction": current.get('winddir16Point', 'N/A'),
                    "pressure": current.get('pressure', 'N/A'),
                    "visibility": current.get('visibility', 'N/A'),
                    "uv_index": current.get('uvIndex', 'N/A')
                },
                "source": "wttr.in (免费天气API)",
                "timestamp": current.get('localObsDateTime', '')
            }
            
            return json.dumps(weather_info, ensure_ascii=False, indent=2)
        else:
            # 如果API调用失败，返回错误信息
            error_info = {
                "error": f"天气API调用失败，状态码: {response.status_code}",
                "location": loc,
                "message": "请检查城市名称是否正确，或稍后重试"
            }
            return json.dumps(error_info, ensure_ascii=False, indent=2)
            
    except requests.exceptions.RequestException as e:
        # 网络请求异常
        error_info = {
            "error": f"网络请求异常: {str(e)}",
            "location": loc,
            "message": "请检查网络连接或稍后重试"
        }
        return json.dumps(error_info, ensure_ascii=False, indent=2)
    except Exception as e:
        # 其他异常
        error_info = {
            "error": f"查询天气时发生错误: {str(e)}",
            "location": loc,
            "message": "请稍后重试或联系技术支持"
        }
        return json.dumps(error_info, ensure_ascii=False, indent=2)

tools = [get_weather]

# 创建模型
model = ChatDeepSeek(model="deepseek-chat")

prompt = """
你是一名乐于助人的智能助手，擅长根据用户的问题选择合适的工具来查询信息并回答。

当用户的问题涉及**天气信息**时，你应优先调用`get_weather`工具，查询用户指定城市的实时天气，并在回答中总结查询结果。

当用户的问题涉及**新闻、事件、实时动态**时，你应优先调用`search_tool`工具，检索相关的最新信息，并在回答中简要概述。

如果问题既包含天气又包含新闻，请先使用`get_weather`查询天气，再使用`search_tool`查询新闻，最后将结果合并后回复用户。

所有回答应使用**简体中文**，条理清晰、简洁友好。

请记住我们之前的对话内容，能够引用和回忆之前讨论过的信息。
"""

# 使用标准的MemorySaver（内存中的checkpointer）
memory = MemorySaver()

# 创建图（带记忆功能）
graph = create_react_agent(model=model, 
                           tools=tools, 
                           prompt=prompt,
                           checkpointer=memory)

print(f"✅ LangGraph已配置内存记忆存储器 + 自定义SQLite持久化")

